// Copyright (C) Kumo inc. and its affiliates.
// Author: Jeff.li lijippy@163.com
// All rights reserved.
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU Affero General Public License as published
// by the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU Affero General Public License for more details.
//
// You should have received a copy of the GNU Affero General Public License
// along with this program.  If not, see <https://www.gnu.org/licenses/>.
//


#include <optional>
#include <pollux/testing/gtest_utils.h>
#include <pollux/testing/functions/prestosql/function_base_test.h>

using namespace kumo::pollux;
using namespace kumo::pollux::test;
using namespace kumo::pollux::functions::test;

namespace {
class ArrayCumSumTest : public FunctionBaseTest {
 protected:
  template <typename T>
  void testArrayCumSum(const VectorPtr& expected, const VectorPtr& input) {
    auto result = evaluate("array_cum_sum(c0)", make_row_vector({input}));
    assertEqualVectors(expected, result);
  }
};

TEST_F(ArrayCumSumTest, bigint) {
  auto input = make_nullable_array_vector<int64_t>(
      {{},
       {10000, 1000010000, 100001000010000},
       {-976543210987654321, std::nullopt, std::numeric_limits<int64_t>::max()},
       {std::nullopt, 100001000010000}});
  auto expected = make_nullable_array_vector<int64_t>(
      {{},
       {10000, 1000020000, 100002000030000},
       {-976543210987654321, std::nullopt, std::nullopt},
       {std::nullopt, std::nullopt}});
  testArrayCumSum<int64_t>(expected, input);
}

TEST_F(ArrayCumSumTest, integer) {
  auto input = make_nullable_array_vector<int32_t>(
      {{},
       {1000, 10001000, 1010001000},
       {-976543210, std::nullopt, std::numeric_limits<int32_t>::max()},
       {std::nullopt, 1010001000}});
  auto expected = make_nullable_array_vector<int32_t>(
      {{},
       {1000, 10002000, 1020003000},
       {-976543210, std::nullopt, std::nullopt},
       {std::nullopt, std::nullopt}});
  testArrayCumSum<int32_t>(expected, input);
}

TEST_F(ArrayCumSumTest, smallint) {
  auto input = make_nullable_array_vector<int16_t>(
      {{},
       {10, 100, 10000},
       {-9876, std::nullopt, std::numeric_limits<int16_t>::max()},
       {std::nullopt, 10000}});
  auto expected = make_nullable_array_vector<int16_t>(
      {{},
       {10, 110, 10110},
       {-9876, std::nullopt, std::nullopt},
       {std::nullopt, std::nullopt}});
  testArrayCumSum<int16_t>(expected, input);
}

TEST_F(ArrayCumSumTest, tinyint) {
  auto input = make_nullable_array_vector<int8_t>(
      {{},
       {1, 2, 4},
       {-99, std::nullopt, std::numeric_limits<int8_t>::max()},
       {std::nullopt, 4}});
  auto expected = make_nullable_array_vector<int8_t>(
      {{},
       {1, 3, 7},
       {-99, std::nullopt, std::nullopt},
       {std::nullopt, std::nullopt}});
  testArrayCumSum<int8_t>(expected, input);
}

TEST_F(ArrayCumSumTest, real) {
  auto input = make_nullable_array_vector<float>(
      {{},
       {1, 2, 3},
       {-9, std::nullopt, std::numeric_limits<float>::max()},
       {std::nullopt, 3}});
  auto expected = make_nullable_array_vector<float>(
      {{},
       {1, 3, 6},
       {-9, std::nullopt, std::nullopt},
       {std::nullopt, std::nullopt}});
  testArrayCumSum<float>(expected, input);
}

TEST_F(ArrayCumSumTest, double) {
  auto input = make_nullable_array_vector<double>(
      {{},
       {1, 2, 3},
       {-9, std::nullopt, std::numeric_limits<double>::max()},
       {std::nullopt, 3}});
  auto expected = make_nullable_array_vector<double>(
      {{},
       {1, 3, 6},
       {-9, std::nullopt, std::nullopt},
       {std::nullopt, std::nullopt}});
  testArrayCumSum<double>(expected, input);
}

TEST_F(ArrayCumSumTest, bigintOverflow) {
  constexpr int64_t kMin = std::numeric_limits<int64_t>::min();
  constexpr int64_t kMax = std::numeric_limits<int64_t>::max();

  // Overflow on 1 + kMax.
  auto input = make_nullable_array_vector<int64_t>({{1, kMax, 2, 3}});
  POLLUX_ASSERT_THROW(
      evaluate("array_cum_sum(c0)", make_row_vector({input})),
      "integer overflow: 1 + 9223372036854775807");

  // Overflow on -1 + kMin.
  input = make_nullable_array_vector<int64_t>({{-1, kMin, -2, -3}});
  POLLUX_ASSERT_THROW(
      evaluate("array_cum_sum(c0)", make_row_vector({input})),
      "integer overflow: -1 + -9223372036854775808");

  // Overflow for array containing consecutive kMin values.
  input = make_nullable_array_vector<int64_t>({{kMax, kMin, kMin}});
  POLLUX_ASSERT_THROW(
      evaluate("array_cum_sum(c0)", make_row_vector({input})),
      "integer overflow: -1 + -9223372036854775808");

  // No overflow when function is called on alternating kMax, kMin values.
  input = make_nullable_array_vector<int64_t>(
      {{kMax, kMin, kMax, kMin},
       {kMax, kMin},
       {kMax, kMin, kMax, kMin, kMax, kMin}});
  auto expected = make_nullable_array_vector<int64_t>(
      {{kMax, -1, kMax - 1, -2},
       {kMax, -1},
       {kMax, -1, kMax - 1, -2, kMax - 2, -3}});
  testArrayCumSum<int64_t>(expected, input);
}

TEST_F(ArrayCumSumTest, doubleLimits) {
  constexpr double kInf = std::numeric_limits<double>::infinity();
  constexpr double kNan = std::numeric_limits<double>::quiet_NaN();
  constexpr double kLowest = std::numeric_limits<double>::lowest();
  constexpr double kMax = std::numeric_limits<double>::max();

  auto input = make_nullable_array_vector<double>(
      {{1, kInf, 123.456},
       {1, kNan, 123.456},
       {1, kLowest, 123.456},
       {1, kMax, 123.456},
       {kLowest, kMax, 123.456},
       {kMax, kLowest, 123.456},
       {1, kMax, kInf},
       {1, kLowest, kInf},
       {kLowest, kNan, 5},
       {kMax, kNan, 5}});
  auto expected = make_nullable_array_vector<double>(
      {{1, kInf, kInf},
       {1, kNan, kNan},
       {1, kLowest, kLowest},
       {1, kMax, kMax},
       {kLowest, 0, 123.456},
       {kMax, 0, 123.456},
       {1, kMax, kInf},
       {1, kLowest, kInf},
       {kLowest, kNan, kNan},
       {kMax, kNan, kNan}});
  testArrayCumSum<double>(expected, input);
}
} // namespace
